#!/usr/bin/python
# This script does...
# Usage: cogCats.py 1.list 2.list 3.list
# 1.list contains COG identities for the genome. 
# Example: 
#    COG0593 ATPase involved in DNA replication initiation
#    COG0112 Glycine/serine hydroxymethyltransferase
# 2.list contains individual COG category with description
# Example:
# [H]     COG0001 Glutamate-1-semialdehyde aminotransferase
# [E]     COG0002 Acetylglutamate semialdehyde dehydrogenase
# [P]     COG0003 Oxyanion-translocating ATPase
# [P]     COG0004 Ammonia permease
# [F]     COG0005 Purine nucleoside phosphorylase
# [E]     COG0006 Xaa-Pro aminopeptidase
# 3.list contains list of COG categories
# Example:
# J       Translation, ribosomal structure and biogenesis
# A       RNA processing and modification
# K       Transcription
# L       Replication, recombination and repair

import sys
import re

def fmt(f):
#    st = '{0:.4}'.format(f)
    st = '{0:.2}'.format(f)
    return st

indiv = re.compile('\[(\w+)\]\t(\w+)\t(.*)')

cogslistfile = open(sys.argv[1], "rU")

indivcogsfile = open(sys.argv[2], "rU")

cogcatfile = open(sys.argv[3], "rU")

cogslist = cogslistfile.readlines()

indivcogs = indivcogsfile.readlines()

cogcats = cogcatfile.readlines()

totalorfs = len(cogslist)

cogcatdict = {}
for line in cogcats:
    i = line.split('\t')
    #cogcatdict[i[0]] = i[1].rstrip()
    cogcatdict[i[0]] = 0 #set the count to zero
cogcatdict['-'] = 0

indivcogdict = {}
for line in indivcogs:
    i = line.strip()
    if indiv.match(i):
        j = indiv.match(i)
        indivcogdict[j.group(2)] = j.group(1)

for line in cogslist:
    i = line.split(' ')
    cogcat = ""
    cog = i[0]
    if cog in indivcogdict:
        cogcat = indivcogdict[cog]
        if len(cogcat) > 1:
            for i in cogcat:
                cogcatdict[i] += 1 #for multi-functional COGs
        else:
            cogcatdict[cogcat] += 1 #for COGs that have only 1 function
    else:
        cogcatdict['-'] += 1

for i, j in cogcatdict.iteritems():
    percent = float(cogcatdict[i])/float(totalorfs) * 100
    print i + "\t" + str(cogcatdict[i]) + "\t" + fmt(percent)

cogslistfile.close()
indivcogsfile.close()
cogcatfile.close()
